package com.insight.common.message.webssh.handler;


import com.alibaba.fastjson.JSON;
import com.jcraft.jsch.Channel;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import com.insight.common.message.websocket.handler.IWebSocketHandler;
import com.insight.common.message.webssh.constant.ConstantPool;
import com.insight.common.message.webssh.pojo.SSHConnectInfo;
import com.insight.common.message.webssh.pojo.WebSSHData;
import com.insight.common.message.webssh.pojo.WebsocketConst;
import org.springframework.stereotype.Component;

import javax.websocket.Session;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * @Description: WebSSH业务逻辑实现
 * @Author: NoCortY
 * @Date: 2020/3/8
 */
@Slf4j
@Component
public class WebSSHHandler implements IWebSocketHandler {

    //存放ssh连接信息的map
    private static final Map<String, SSHConnectInfo> sshMap = new ConcurrentHashMap<>();

    //线程池
    private static final ExecutorService executorService = Executors.newCachedThreadPool();

    @Override
    public boolean support(String support) {
        return StringUtils.equals(support, WebsocketConst.SSH);
    }

    @Override
    public void onOpenHandler(String code, Session session) {
        SSHConnectInfo sshConnectInfo = new SSHConnectInfo();
        //将这个ssh连接信息放入map中
        sshMap.put(session.getId(), sshConnectInfo);
    }

    @Override
    public void onMessageHandler(String buffer, Session session) {
        WebSSHData webSSHData = JSON.parseObject(buffer, WebSSHData.class);
        String sessionId = session.getId();
        SSHConnectInfo sshConnectInfo = sshMap.get(sessionId);
        switch (webSSHData.getOperate()) {
            case ConstantPool.WEBSOCKET_OPERATE_CONNECT:
                //启动线程异步处理
                executorService.execute(() -> {
                    try {
                        connectToSSH(sshConnectInfo, webSSHData, session);
                    } catch (JSchException | IOException e) {
                        exceptionHandle(session, e);
                        //closeSSHConnect(session);
                    }
                });
                break;
            case ConstantPool.WEBSOCKET_OPERATE_COMMAND:
                String command = webSSHData.getCommand();
                try {
                    transToSSH(sshConnectInfo.getChannel(), command);
                } catch (IOException e) {
                    exceptionHandle(session, e);
                    //closeSSHConnect(session);
                }
                break;
            default:
                closeSSHConnect(session);
        }
    }

    private void exceptionHandle(Session session, Exception e) {
        log.error("webssh连接异常");
        log.error("异常信息:{}", e.getMessage());
        synchronized (session) {
            try {
                String message = e.getMessage() != null ? e.getMessage() : "未知异常";
                sendMessage(session, ("系统异常：" + message).getBytes(StandardCharsets.UTF_8));
                sendMessage(session, "\r\n".getBytes(StandardCharsets.UTF_8));
                session.close();
            } catch (IOException ex) {
                log.error("关闭WebSSH会话时发生IO异常", ex);
                throw new IllegalArgumentException("WebSSH异常处理失败", ex);
            }
        }
    }

    @Override
    public void
    onCloseHandler(Session session) {
        closeSSHConnect(session);
    }

    private void closeSSHConnect(Session session) {
        String sessionId = session.getId();
        SSHConnectInfo sshConnectInfo = sshMap.get(sessionId);
        if (sshConnectInfo != null) {
            //断开连接
            if (sshConnectInfo.getChannel() != null) {
                sshConnectInfo.getChannel().disconnect();
            }
            if (sshConnectInfo.getSshSession() != null) {
                sshConnectInfo.getSshSession().disconnect();
            }
            //map中移除
            sshMap.remove(sessionId);
        }
    }

    /**
     * @Description: 使用jsch连接终端
     * @Param: [cloudSSH, webSSHData, webSocketSession]
     * @return: void
     * @Author: NoCortY
     * @Date: 2020/3/7
     */
    private void connectToSSH(SSHConnectInfo sshConnectInfo, WebSSHData webSSHData, Session webSocketSession) throws JSchException, IOException {
        Properties config = new Properties();
        config.put("StrictHostKeyChecking", "no");
        //获取jsch的会话
        com.jcraft.jsch.Session session = null;
        if (sshConnectInfo.getSshSession() != null) {
            session = sshConnectInfo.getSshSession();
        } else {
            JSch jSch = new JSch();
            session = jSch
                    .getSession(webSSHData.getUsername(), webSSHData.getHost(), webSSHData.getPort());
            //设置session
            sshConnectInfo.setSshSession(session);
        }

        session.setConfig(config);
        //设置密码
        session.setPassword(webSSHData.getPassword());
        //连接  超时时间30s
        session.connect(30000);

        //开启shell通道
        Channel channel = session.openChannel("shell");

        //通道连接 超时时间3s
        channel.connect(3000);

        //设置channel
        sshConnectInfo.setChannel(channel);

        //转发消息
        //transToSSH(channel, "\r");

        //读取终端返回的信息流
        try (InputStream inputStream = channel.getInputStream()) {
            //循环读取
            byte[] buffer = new byte[1024];
            int i = 0;
            //如果没有数据来，线程会一直阻塞在这个地方等待数据。
            while ((i = inputStream.read(buffer)) != -1) {
                synchronized (webSocketSession) {
                    sendMessage(webSocketSession, Arrays.copyOfRange(buffer, 0, i));
                }
            }
        }
    }

    /**
     * @Description: 将消息转发到终端
     * @Param: [channel, data]
     * @return: void
     * @Author: NoCortY
     * @Date: 2020/3/7
     */
    private void transToSSH(Channel channel, String command) throws IOException {
        if (channel != null) {
            OutputStream outputStream = channel.getOutputStream();
            outputStream.write(command.getBytes());
            outputStream.flush();
        }
    }

    private void sendMessage(Session session, byte[] buffer) throws IOException {
        session.getBasicRemote().sendText(new String(buffer, StandardCharsets.UTF_8));
    }
}
